#include "conv2d.h"
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cstdio>
#include <cassert>
#include "core.h"


namespace uzu
{
#define MAX_BUFFER_LENGTH (SQUARE_BLOCK_SIZE + MAX_CONV2D_KERNEL_SIZE)

    __global__ void conv2d_kernel(float* source, float* filter, float* dest, 
        int width, int height, int filter_width, int filter_height)
    {
        __shared__ float buffer[MAX_BUFFER_LENGTH * MAX_BUFFER_LENGTH];
        __shared__ float filter_buffer[MAX_CONV2D_KERNEL_SIZE * MAX_CONV2D_KERNEL_SIZE];

        int half_filter_x = filter_width / 2;
        int half_filter_y = filter_height / 2;

        int real_block_size_x = width - blockIdx.x * blockDim.x;
        if (real_block_size_x >= blockDim.x) real_block_size_x = blockDim.x;
        int real_block_size_y = height - blockIdx.y * blockDim.y;
        if (real_block_size_y >= blockDim.y) real_block_size_y = blockDim.y;

        // copy filter to filter_buffer
        if (threadIdx.x < filter_width && threadIdx.y < filter_height)
        {
            int filter_buffer_index = INDEX_2D(threadIdx.x, threadIdx.y, MAX_CONV2D_KERNEL_SIZE, MAX_CONV2D_KERNEL_SIZE);
            int filter_index = INDEX_2D(threadIdx.x, threadIdx.y, filter_width, filter_height);
            filter_buffer[filter_buffer_index] = filter[filter_index];
        }

        // copy source to buffer
        int start_x = blockIdx.x * blockDim.x - half_filter_x;
        int end_x = blockIdx.x * blockDim.x + real_block_size_x + filter_width - half_filter_x - 1;
        int nrepeats_x = (end_x - start_x + SQUARE_BLOCK_SIZE - 1) / SQUARE_BLOCK_SIZE;
        int start_y = blockIdx.y * blockDim.y - half_filter_y;
        int end_y = blockIdx.y * blockDim.y + real_block_size_y + filter_height - half_filter_y - 1;
        int nrepeats_y = (end_y - start_y + SQUARE_BLOCK_SIZE - 1) / SQUARE_BLOCK_SIZE;

        for (int rx = 0; rx < nrepeats_x; ++rx)
        {
            for (int ry = 0; ry < nrepeats_y; ++ry)
            {
                int buffer_x = rx * SQUARE_BLOCK_SIZE + threadIdx.x;
                int buffer_y = ry * SQUARE_BLOCK_SIZE + threadIdx.y;
                int source_x = start_x + buffer_x;
                int source_y = start_y + buffer_y;
                if (buffer_x < real_block_size_x + filter_width - 1 && 
                    buffer_y < real_block_size_y + filter_height - 1)
                {
                    if (source_x < 0 || source_y < 0)
                        buffer[INDEX_2D(buffer_x, buffer_y, MAX_BUFFER_LENGTH, MAX_BUFFER_LENGTH)] = 0.0f;
                    else if (source_x >= width || source_y >= height || source_x >= end_x || source_y >= end_y)
                        buffer[INDEX_2D(buffer_x, buffer_y, MAX_BUFFER_LENGTH, MAX_BUFFER_LENGTH)] = 0.0f;
                    else
                        buffer[INDEX_2D(buffer_x, buffer_y, MAX_BUFFER_LENGTH, MAX_BUFFER_LENGTH)] = source[INDEX_2D(source_x, source_y, width, height)];
                }
            }
        }

        __syncthreads();

        // do calculation
        int index_x = threadIdx.x + blockIdx.x * blockDim.x;
        int index_y = threadIdx.y + blockIdx.y * blockDim.y;
        if (index_x < width && index_y < height)
        {
            float sum = 0.0f;
            for (int i = 0; i < filter_width; ++i)
            {
                for (int j = 0; j < filter_height; ++j)
                {
                    int buffer_index = INDEX_2D(threadIdx.x + i, threadIdx.y + j, 
                        MAX_BUFFER_LENGTH, MAX_BUFFER_LENGTH);
                    int filter_index = INDEX_2D(i, j, MAX_CONV2D_KERNEL_SIZE, MAX_CONV2D_KERNEL_SIZE);
                    sum += buffer[buffer_index] * filter_buffer[filter_index];
                }
            }
            dest[INDEX_2D(index_x, index_y, width, height)] = sum;
        }
    }

    void Conv2d(Tensor& source, Tensor& filter, Tensor& dest)
    {
        assert(source.NDim() == 2);
        assert(filter.NDim() == 2);
        auto source_shape = source.Shape();
        auto filter_shape = filter.Shape();
        int width = source_shape[0];
        int height = source_shape[1];
        int filter_width = filter_shape[0];
        int filter_height = filter_shape[1];

        assert(filter_width < MAX_CONV2D_KERNEL_SIZE);
        assert(filter_height < MAX_CONV2D_KERNEL_SIZE);

        dest.Create(source_shape, DeviceType::CUDA);
        
        source.ToDevice();
        filter.ToDevice();

        // launcha kernel;
        dim3 block_size(SQUARE_BLOCK_SIZE, SQUARE_BLOCK_SIZE, 1);
        int blocks_x = (width + SQUARE_BLOCK_SIZE - 1) / SQUARE_BLOCK_SIZE;
        int blocks_y = (height + SQUARE_BLOCK_SIZE - 1) / SQUARE_BLOCK_SIZE;
        dim3 grid_size(blocks_x, blocks_y, 1);
        conv2d_kernel<<<grid_size, block_size>>>(
            source.DataGpu(), filter.DataGpu(), dest.DataGpu(), 
            width, height, filter_width, filter_height);
    }
}
